﻿#define _WINSOCK_DEPRECATED_NO_WARNINGS
#define _CRT_SECURE_NO_WARNINGS
#include <winsock2.h>
#include <windows.h>
#include <process.h>
#include <stdio.h>
#include <stdarg.h>
#include <time.h>
#include <stdlib.h>
#include <string.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")

// 添加getopt实现
extern char* optarg;
extern int optind, opterr, optopt;

int getopt(int argc, char* const argv[], const char* optstring) {
    static int optpos = 1;
    const char* p;  // 修改为const char*类型

    if (optind >= argc || argv[optind][0] != '-' || argv[optind][1] == '\0') {
        return -1;
    }

    if (strcmp(argv[optind], "--") == 0) {
        optind++;
        return -1;
    }

    optopt = argv[optind][optpos];
    p = strchr(optstring, optopt);  // 匹配strchr的返回类型

    if (p == NULL) {
        if (opterr) {
            fprintf(stderr, "%s: illegal option -- %c\n", argv[0], optopt);
        }
        if (argv[optind][++optpos] == '\0') {
            optind++;
            optpos = 1;
        }
        return '?';
    }

    if (p[1] == ':') {
        if (argv[optind][optpos + 1] != '\0') {
            optarg = &argv[optind][optpos + 1];
            optind++;
            optpos = 1;
        }
        else if (optind + 1 < argc) {
            optarg = argv[++optind];
            optind++;
            optpos = 1;
        }
        else {
            if (opterr) {
                fprintf(stderr, "%s: option requires an argument -- %c\n", argv[0], optopt);
            }
            optind++;
            optpos = 1;
            return '?';
        }
    }
    else {
        if (argv[optind][++optpos] == '\0') {
            optind++;
            optpos = 1;
        }
    }

    return optopt;
}

char* optarg = NULL;
int optind = 1;
int opterr = 1;
int optopt = 0;

#define BUFSIZE 65536
#define IPSIZE 4
#define ARRAY_SIZE(x) (sizeof(x) / sizeof(x[0]))

unsigned short int port = 1080;
FILE* log_file = stdout;
int auth_type = 0;
char* arg_username = NULL;
char* arg_password = NULL;
CRITICAL_SECTION lock;

enum socks {
    RESERVED = 0x00,
    VERSION4 = 0x04,
    VERSION5 = 0x05
};

enum socks_auth_methods {
    NOAUTH = 0x00,
    USERPASS = 0x02,
    NOMETHOD = 0xff
};

enum socks_auth_userpass {
    AUTH_OK = 0x00,
    AUTH_VERSION = 0x01,
    AUTH_FAIL = 0xff
};

enum socks_command {
    CONNECT = 0x01
};

enum socks_command_type {
    IP = 0x01,
    DOMAIN_NAME = 0x03  // 修改为DOMAIN_NAME避免与预定义宏冲突
};

enum socks_status {
    OK = 0x00,
    FAILED = 0x05
};

void log_message(const char* message, ...) {
    char vbuffer[255];
    va_list args;
    va_start(args, message);
    vsnprintf(vbuffer, ARRAY_SIZE(vbuffer), message, args);
    va_end(args);

    time_t now;
    time(&now);
    char* date = ctime(&now);
    date[strlen(date) - 1] = '\0';

    EnterCriticalSection(&lock);
    if (WSAGetLastError() != 0) {
        fprintf(log_file, "[%s][%lu] Critical: %s - %s\n", date, GetCurrentThreadId(), vbuffer, gai_strerror(WSAGetLastError()));
        WSASetLastError(0);
    }
    else {
        fprintf(log_file, "[%s][%lu] Info: %s\n", date, GetCurrentThreadId(), vbuffer);
    }
    fflush(log_file);
    LeaveCriticalSection(&lock);
}

int recvn(SOCKET fd, void* buf, int n) {
    int nread, left = n;
    while (left > 0) {
        nread = recv(fd, (char*)buf + (n - left), left, 0);
        if (nread == SOCKET_ERROR) {
            int error = WSAGetLastError();
            if (error == WSAEINTR || error == WSAEWOULDBLOCK) {
                continue;
            }
            log_message("recv failed: %d", error);
            return -1;
        }
        else if (nread == 0) {
            return 0;
        }
        left -= nread;
    }
    return n;
}

int sendn(SOCKET fd, void* buf, int n) {
    int nwrite, left = n;
    while (left > 0) {
        nwrite = send(fd, (char*)buf + (n - left), left, 0);
        if (nwrite == SOCKET_ERROR) {
            int error = WSAGetLastError();
            if (error == WSAEINTR || error == WSAEWOULDBLOCK) {
                continue;
            }
            log_message("send failed: %d", error);
            return -1;
        }
        else if (nwrite == 0) {
            return 0;
        }
        left -= nwrite;
    }
    return n;
}

void app_thread_exit(int ret, SOCKET fd) {
    closesocket(fd);
    _endthreadex(ret);
}

int app_connect(int type, void* buf, unsigned short int portnum) {
    SOCKET fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (fd == INVALID_SOCKET) {
        log_message("socket failed: %d", WSAGetLastError());
        return -1;
    }

    struct sockaddr_in remote;
    char address[16];

    if (type == IP) {
        char* ip = (char*)buf;
        snprintf(address, ARRAY_SIZE(address), "%hhu.%hhu.%hhu.%hhu", ip[0], ip[1], ip[2], ip[3]);
        memset(&remote, 0, sizeof(remote));
        remote.sin_family = AF_INET;
        remote.sin_addr.s_addr = inet_addr(address);
        remote.sin_port = htons(portnum);

        if (connect(fd, (struct sockaddr*)&remote, sizeof(remote)) == SOCKET_ERROR) {
            log_message("connect failed: %d", WSAGetLastError());
            closesocket(fd);
            return -1;
        }
    }
    else if (type == DOMAIN_NAME) {  // 使用修改后的枚举值
        char portaddr[6];
        struct addrinfo* res, hints;
        snprintf(portaddr, ARRAY_SIZE(portaddr), "%d", portnum);
        log_message("getaddrinfo: %s %s", (const char*)buf, portaddr);

        memset(&hints, 0, sizeof(hints));
        hints.ai_family = AF_INET;
        hints.ai_socktype = SOCK_STREAM;

        int ret = getaddrinfo((const char*)buf, portaddr, &hints, &res);
        if (ret != 0) {
            log_message("getaddrinfo failed: %d", ret);
            closesocket(fd);
            return -1;
        }

        struct addrinfo* r;
        for (r = res; r != NULL; r = r->ai_next) {
            if (connect(fd, r->ai_addr, r->ai_addrlen) == 0) {
                freeaddrinfo(res);
                return fd;
            }
        }

        freeaddrinfo(res);
        closesocket(fd);
        return -1;
    }

    return fd;
}

int socks_invitation(SOCKET fd, int* version) {
    char init[2];
    int nread = recvn(fd, (void*)init, ARRAY_SIZE(init));
    if (nread == 2 && init[0] != VERSION5 && init[0] != VERSION4) {
        log_message("They send us %hhX %hhX", init[0], init[1]);
        log_message("Incompatible version!");
        app_thread_exit(0, fd);
    }
    log_message("Initial %hhX %hhX", init[0], init[1]);
    *version = init[0];
    return init[1];
}

char* socks5_auth_get_user(SOCKET fd) {
    unsigned char size;
    recvn(fd, (void*)&size, sizeof(size));

    char* user = (char*)malloc(sizeof(char) * size + 1);
    if (user == NULL) {
        log_message("Memory allocation failed for username");
        app_thread_exit(1, fd);
    }
    recvn(fd, (void*)user, (int)size);
    user[size] = 0;

    return user;
}

char* socks5_auth_get_pass(SOCKET fd) {
    unsigned char size;
    recvn(fd, (void*)&size, sizeof(size));

    char* pass = (char*)malloc(sizeof(char) * size + 1);
    if (pass == NULL) {
        log_message("Memory allocation failed for password");
        app_thread_exit(1, fd);
    }
    recvn(fd, (void*)pass, (int)size);
    pass[size] = 0;

    return pass;
}

int socks5_auth_userpass(SOCKET fd) {
    char answer[2] = { VERSION5, USERPASS };
    sendn(fd, (void*)answer, ARRAY_SIZE(answer));
    char resp;
    recvn(fd, (void*)&resp, sizeof(resp));
    log_message("auth %hhX", resp);
    char* username = socks5_auth_get_user(fd);
    char* password = socks5_auth_get_pass(fd);
    log_message("l: %s p: %s", username, password);
    int result = 1;
    if (arg_username && arg_password &&
        strcmp(arg_username, username) == 0 &&
        strcmp(arg_password, password) == 0) {
        char answer[2] = { AUTH_VERSION, AUTH_OK };
        sendn(fd, (void*)answer, ARRAY_SIZE(answer));
        result = 0;
    }
    else {
        char answer[2] = { AUTH_VERSION, AUTH_FAIL };
        sendn(fd, (void*)answer, ARRAY_SIZE(answer));
    }
    free(username);
    free(password);
    return result;
}

int socks5_auth_noauth(SOCKET fd) {
    char answer[2] = { VERSION5, NOAUTH };
    sendn(fd, (void*)answer, ARRAY_SIZE(answer));
    return 0;
}

void socks5_auth_notsupported(SOCKET fd) {
    char answer[2] = { VERSION5, NOMETHOD };
    sendn(fd, (void*)answer, ARRAY_SIZE(answer));
}

void socks5_auth(SOCKET fd, int methods_count) {
    int supported = 0;
    int num = methods_count;
    for (int i = 0; i < num; i++) {
        char type;
        recvn(fd, (void*)&type, 1);
        log_message("Method AUTH %hhX", type);
        if (type == auth_type) {
            supported = 1;
        }
    }
    if (supported == 0) {
        socks5_auth_notsupported(fd);
        app_thread_exit(1, fd);
    }
    int ret = 0;
    switch (auth_type) {
    case NOAUTH:
        ret = socks5_auth_noauth(fd);
        break;
    case USERPASS:
        ret = socks5_auth_userpass(fd);
        break;
    }
    if (ret == 0) {
        return;
    }
    else {
        app_thread_exit(1, fd);
    }
}

int socks5_command(SOCKET fd) {
    char command[4];
    recvn(fd, (void*)command, ARRAY_SIZE(command));
    log_message("Command %hhX %hhX %hhX %hhX", command[0], command[1], command[2], command[3]);
    return command[3];
}

unsigned short int socks_read_port(SOCKET fd) {
    unsigned short int p;
    recvn(fd, (void*)&p, sizeof(p));
    log_message("Port %hu", ntohs(p));
    return p;
}

char* socks_ip_read(SOCKET fd) {
    char* ip = (char*)malloc(sizeof(char) * IPSIZE);
    if (ip == NULL) {
        log_message("Memory allocation failed for IP");
        app_thread_exit(1, fd);
    }
    recvn(fd, (void*)ip, IPSIZE);
    log_message("IP %hhu.%hhu.%hhu.%hhu", ip[0], ip[1], ip[2], ip[3]);
    return ip;
}

void socks5_ip_send_response(SOCKET fd, char* ip, unsigned short int port) {
    char response[4] = { VERSION5, OK, RESERVED, IP };
    sendn(fd, (void*)response, ARRAY_SIZE(response));
    sendn(fd, (void*)ip, IPSIZE);
    sendn(fd, (void*)&port, sizeof(port));
}

char* socks5_domain_read(SOCKET fd, unsigned char* size) {
    unsigned char s;
    recvn(fd, (void*)&s, sizeof(s));
    char* address = (char*)malloc((sizeof(char) * s) + 1);
    if (address == NULL) {
        log_message("Memory allocation failed for domain");
        app_thread_exit(1, fd);
    }
    recvn(fd, (void*)address, (int)s);
    address[s] = 0;
    log_message("Address %s", address);
    *size = s;
    return address;
}

void socks5_domain_send_response(SOCKET fd, char* domain, unsigned char size, unsigned short int port) {
    char response[4] = { VERSION5, OK, RESERVED, DOMAIN_NAME };  // 使用修改后的枚举值
    sendn(fd, (void*)response, ARRAY_SIZE(response));
    sendn(fd, (void*)&size, sizeof(size));
    sendn(fd, (void*)domain, size * sizeof(char));
    sendn(fd, (void*)&port, sizeof(port));
}

int socks4_is_4a(char* ip) {
    return (ip[0] == 0 && ip[1] == 0 && ip[2] == 0 && ip[3] != 0);
}

int socks4_read_nstring(SOCKET fd, char* buf, int size) {
    char sym = 0;
    int nread = 0;
    int i = 0;

    while (i < size) {
        nread = recv(fd, &sym, sizeof(char), 0);

        if (nread <= 0) {
            break;
        }
        else {
            buf[i] = sym;
            i++;
        }

        if (sym == 0) {
            break;
        }
    }

    return i;
}

void socks4_send_response(SOCKET fd, int status) {
    char resp[8] = { 0x00, (char)status, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };
    sendn(fd, (void*)resp, ARRAY_SIZE(resp));
}

void app_socket_pipe(SOCKET fd0, SOCKET fd1) {
    int maxfd, ret;
    fd_set rd_set;
    size_t nread;
    char buffer_r[BUFSIZE];

    log_message("Connecting two sockets");

    maxfd = (fd0 > fd1) ? fd0 : fd1;
    while (1) {
        FD_ZERO(&rd_set);
        FD_SET(fd0, &rd_set);
        FD_SET(fd1, &rd_set);
        ret = select(maxfd + 1, &rd_set, NULL, NULL, NULL);

        if (ret < 0) {
            int error = WSAGetLastError();
            if (error == WSAEINTR) {
                continue;
            }
            log_message("select failed: %d", error);
            break;
        }

        if (FD_ISSET(fd0, &rd_set)) {
            nread = recv(fd0, buffer_r, BUFSIZE, 0);
            if (nread <= 0)
                break;
            send(fd1, (const char*)buffer_r, nread, 0);  // 显式类型转换
        }

        if (FD_ISSET(fd1, &rd_set)) {
            nread = recv(fd1, buffer_r, BUFSIZE, 0);
            if (nread <= 0)
                break;
            send(fd0, (const char*)buffer_r, nread, 0);  // 显式类型转换
        }
    }
}

unsigned int __stdcall app_thread_process(void* arg) {
    SOCKET net_fd = *(SOCKET*)arg;
    free(arg);

    int version = 0;
    SOCKET inet_fd = INVALID_SOCKET;
    char methods = socks_invitation(net_fd, &version);

    switch (version) {
    case VERSION5: {
        socks5_auth(net_fd, methods);
        int command = socks5_command(net_fd);

        if (command == IP) {
            char* ip = socks_ip_read(net_fd);
            unsigned short int p = socks_read_port(net_fd);

            inet_fd = app_connect(IP, (void*)ip, ntohs(p));
            if (inet_fd == INVALID_SOCKET) {
                app_thread_exit(1, net_fd);
            }
            socks5_ip_send_response(net_fd, ip, p);
            free(ip);
        }
        else if (command == DOMAIN_NAME) {  // 使用修改后的枚举值
            unsigned char size;
            char* address = socks5_domain_read(net_fd, &size);
            unsigned short int p = socks_read_port(net_fd);

            inet_fd = app_connect(DOMAIN_NAME, (void*)address, ntohs(p));
            if (inet_fd == INVALID_SOCKET) {
                app_thread_exit(1, net_fd);
            }
            socks5_domain_send_response(net_fd, address, size, p);
            free(address);
        }
        else {
            app_thread_exit(1, net_fd);
        }
        break;
    }
    case VERSION4: {
        if (methods == 1) {
            char ident[255];
            unsigned short int p = socks_read_port(net_fd);
            char* ip = socks_ip_read(net_fd);
            socks4_read_nstring(net_fd, ident, sizeof(ident));

            if (socks4_is_4a(ip)) {
                char domain[255];
                socks4_read_nstring(net_fd, domain, sizeof(domain));
                log_message("Socks4A: ident:%s; domain:%s;", ident, domain);
                inet_fd = app_connect(DOMAIN_NAME, (void*)domain, ntohs(p));  // 使用修改后的枚举值
            }
            else {
                log_message("Socks4: connect by ip & port");
                inet_fd = app_connect(IP, (void*)ip, ntohs(p));
            }

            if (inet_fd != INVALID_SOCKET) {
                socks4_send_response(net_fd, 0x5a);
            }
            else {
                socks4_send_response(net_fd, 0x5b);
                free(ip);
                app_thread_exit(1, net_fd);
            }

            free(ip);
        }
        else {
            log_message("Unsupported mode");
        }
        break;
    }
    }

    app_socket_pipe(inet_fd, net_fd);
    closesocket(inet_fd);
    app_thread_exit(0, net_fd);
    return 0;
}

int app_loop() {
    SOCKET sock_fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock_fd == INVALID_SOCKET) {
        log_message("socket failed: %d", WSAGetLastError());
        exit(1);
    }

    int optval = 1;
    if (setsockopt(sock_fd, SOL_SOCKET, SO_REUSEADDR, (char*)&optval, sizeof(optval)) == SOCKET_ERROR) {
        log_message("setsockopt failed: %d", WSAGetLastError());
        closesocket(sock_fd);
        exit(1);
    }

    struct sockaddr_in local;
    memset(&local, 0, sizeof(local));
    local.sin_family = AF_INET;
    local.sin_addr.s_addr = htonl(INADDR_ANY);
    local.sin_port = htons(port);

    if (bind(sock_fd, (struct sockaddr*)&local, sizeof(local)) == SOCKET_ERROR) {
        log_message("bind failed: %d", WSAGetLastError());
        closesocket(sock_fd);
        exit(1);
    }

    if (listen(sock_fd, 25) == SOCKET_ERROR) {
        log_message("listen failed: %d", WSAGetLastError());
        closesocket(sock_fd);
        exit(1);
    }

    log_message("Listening port %d...", port);

    while (1) {
        struct sockaddr_in remote;
        int remotelen = sizeof(remote);
        SOCKET net_fd = accept(sock_fd, (struct sockaddr*)&remote, &remotelen);
        if (net_fd == INVALID_SOCKET) {
            log_message("accept failed: %d", WSAGetLastError());
            continue;
        }

        int one = 1;
        if (setsockopt(net_fd, IPPROTO_TCP, TCP_NODELAY, (char*)&one, sizeof(one)) == SOCKET_ERROR) {
            log_message("setsockopt TCP_NODELAY failed: %d", WSAGetLastError());
        }

        SOCKET* pNetFd = (SOCKET*)malloc(sizeof(SOCKET));
        if (pNetFd == NULL) {
            log_message("Memory allocation failed for client socket");
            closesocket(net_fd);
            continue;
        }
        *pNetFd = net_fd;
        HANDLE worker = (HANDLE)_beginthreadex(NULL, 0, app_thread_process, pNetFd, 0, NULL);
        if (worker == 0) {
            log_message("Create thread failed");
            free(pNetFd);
            closesocket(net_fd);
        }
        else {
            CloseHandle(worker);
        }
    }

    closesocket(sock_fd);
    return 0;
}

void usage(char* app) {
    printf("USAGE: %s [-h][-n PORT][-a AUTHTYPE][-u USERNAME][-p PASSWORD][-l LOGFILE]\n", app);
    printf("AUTHTYPE: 0 for NOAUTH, 2 for USERPASS\n");
    printf("By default: port is 1080, authtype is no auth, logfile is stdout\n");
    exit(1);
}

int main(int argc, char* argv[]) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        printf("WSAStartup failed: %d\n", WSAGetLastError());
        return 1;
    }

    InitializeCriticalSection(&lock);
    log_file = stdout;
    auth_type = NOAUTH;

    // 初始化为NULL，在main中分配内存
    arg_username = NULL;
    arg_password = NULL;

    int ret;
    while ((ret = getopt(argc, argv, "n:u:p:l:a:h")) != -1) {
        switch (ret) {
        case 'n':
            port = atoi(optarg) & 0xffff;
            break;
        case 'u':
            if (arg_username) free(arg_username);
            arg_username = _strdup(optarg);
            if (!arg_username) {
                printf("Memory allocation failed for username\n");
                return 1;
            }
            break;
        case 'p':
            if (arg_password) free(arg_password);
            arg_password = _strdup(optarg);
            if (!arg_password) {
                printf("Memory allocation failed for password\n");
                return 1;
            }
            break;
        case 'l':
            log_file = fopen(optarg, "a");
            if (log_file == NULL) {
                printf("Failed to open log file: %s\n", optarg);
                return 1;
            }
            break;
        case 'a':
            auth_type = atoi(optarg);
            break;
        case 'h':
        default:
            usage(argv[0]);
        }
    }

    // 使用_strdup替代strdup，符合Windows安全标准
    if (!arg_username) {
        arg_username = _strdup("user");
        if (!arg_username) {
            printf("Memory allocation failed for default username\n");
            return 1;
        }
    }

    if (!arg_password) {
        arg_password = _strdup("pass");
        if (!arg_password) {
            printf("Memory allocation failed for default password\n");
            return 1;
        }
    }

    log_message("Starting with authtype %X", auth_type);
    if (auth_type != NOAUTH) {
        log_message("Username is %s, password is %s", arg_username, arg_password);
    }

    int result = app_loop();

    // 确保释放所有分配的内存
    if (arg_username) free(arg_username);
    if (arg_password) free(arg_password);
    if (log_file != stdout) fclose(log_file);
    DeleteCriticalSection(&lock);
    WSACleanup();
    return result;
}